#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""目标检测引擎抽象基类

定义统一的检测引擎接口，为不同的检测引擎实现提供标准
"""

from abc import ABC, abstractmethod
from typing import List, Dict, Optional, Tuple, Any
import numpy as np
import time
import logging

# 设置日志记录器
logger = logging.getLogger(__name__)

class BaseDetectionEngine(ABC):
    """目标检测引擎抽象基类"""
    
    def __init__(self, device: str = 'auto', **kwargs):
        """初始化检测引擎
        
        Args:
            device: 运行设备，'auto'、'cuda'、'cpu'
            **kwargs: 其他配置参数
        """
        self.device = device
        self.is_initialized = False
        self.last_detection_time = 0
        self.detection_count = 0
        self.fps = 0.0
        self.total_inference_time = 0.0
        self.avg_inference_time = 0.0
        
    @abstractmethod
    def initialize(self) -> bool:
        """初始化检测引擎
        
        Returns:
            bool: 是否初始化成功
        """
        pass
    
    @abstractmethod
    def detect(self, frame: np.ndarray, **kwargs) -> List[Dict[str, Any]]:
        """执行目标检测
        
        Args:
            frame: 输入图像帧
            **kwargs: 额外检测参数
                - confidence_threshold: 置信度阈值
                - classes: 要检测的类别列表
                - agnostic_nms: 是否使用类别无关的NMS
                
        Returns:
            List[Dict]: 检测结果列表，每个结果包含bbox、置信度、类别等信息
        """
        pass
    
    def batch_detect(self, frames: List[np.ndarray], **kwargs) -> List[List[Dict[str, Any]]]:
        """批量执行目标检测
        
        Args:
            frames: 输入图像帧列表
            **kwargs: 额外检测参数
                
        Returns:
            List[List[Dict]]: 每帧的检测结果列表
        """
        results = []
        for frame in frames:
            results.append(self.detect(frame, **kwargs))
        return results
    
    def track(self, detections: List[Dict[str, Any]], frame: np.ndarray) -> List[Dict[str, Any]]:
        """跟踪检测到的目标
        
        Args:
            detections: 检测结果列表
            frame: 当前图像帧
            
        Returns:
            List[Dict]: 跟踪结果列表，包含跟踪ID和轨迹信息
        """
        # 基类提供空实现，具体实现由子类提供
        return detections
    
    def detect_and_track(self, frame: np.ndarray, **kwargs) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
        """执行检测和跟踪
        
        Args:
            frame: 输入图像帧
            **kwargs: 额外检测参数
                
        Returns:
            Tuple[List[Dict], List[Dict]]: 检测结果和跟踪结果
        """
        detections = self.detect(frame, **kwargs)
        tracked_objects = self.track(detections, frame)
        return detections, tracked_objects
    
    def visualize_results(self, frame: np.ndarray, detections: List[Dict[str, Any]], 
                         show_confidence: bool = True, show_tracks: bool = False, 
                         tracked_objects: Optional[List[Dict[str, Any]]] = None) -> np.ndarray:
        """可视化检测结果
        
        Args:
            frame: 输入图像帧
            detections: 检测结果列表
            show_confidence: 是否显示置信度
            show_tracks: 是否显示跟踪轨迹
            tracked_objects: 跟踪结果列表
            
        Returns:
            np.ndarray: 可视化后的图像
        """
        # 基类提供基本实现，子类可以重写以提供更复杂的可视化
        import cv2
        
        result_frame = frame.copy()
        
        # 为不同类别定义颜色
        colors = {
            'person': (0, 0, 255),      # 红色
            'car': (0, 255, 0),        # 绿色
            'bicycle': (255, 165, 0),   # 橙色
            'motorcycle': (255, 0, 255), # 紫色
            'truck': (0, 255, 255),     # 青色
            'bus': (128, 0, 128),       # 紫色
            'default': (128, 128, 128)  # 灰色
        }
        
        # 绘制检测框
        for det in detections:
            x1, y1, x2, y2 = det.get('x1', 0), det.get('y1', 0), det.get('x2', 0), det.get('y2', 0)
            class_name = det.get('class_name', 'unknown')
            confidence = det.get('confidence', 0.0)
            
            # 获取颜色
            color = colors.get(class_name.lower(), colors['default'])
            
            # 绘制边界框
            cv2.rectangle(result_frame, (x1, y1), (x2, y2), color, 2)
            
            # 准备标签文本
            label_parts = [class_name]
            if show_confidence:
                label_parts.append(f"{confidence:.2f}")
            
            label = ' '.join(label_parts)
            
            # 计算标签位置
            label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)
            label_y = max(y1 - 10, label_size[1])
            
            # 绘制标签背景
            cv2.rectangle(
                result_frame,
                (x1, label_y - label_size[1] - 5),
                (x1 + label_size[0] + 10, label_y + 5),
                color,
                -1
            )
            
            # 绘制标签文本
            cv2.putText(
                result_frame,
                label,
                (x1 + 5, label_y),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (255, 255, 255),
                2
            )
        
        # 如果需要显示跟踪轨迹
        if show_tracks and tracked_objects:
            for track_obj in tracked_objects:
                trajectory = track_obj.get('trajectory', [])
                if len(trajectory) > 1:
                    # 绘制轨迹线
                    for i in range(1, len(trajectory)):
                        cv2.line(
                            result_frame,
                            (int(trajectory[i-1][0]), int(trajectory[i-1][1])),
                            (int(trajectory[i][0]), int(trajectory[i][1])),
                            colors.get(track_obj.get('class_name', 'default').lower(), colors['default']),
                            2
                        )
        
        return result_frame
    
    def get_performance_metrics(self) -> Dict[str, Any]:
        """获取检测引擎性能指标
        
        Returns:
            Dict: 性能指标字典
        """
        return {
            'initialized': self.is_initialized,
            'device': self.device,
            'detection_count': self.detection_count,
            'fps': round(self.fps, 2),
            'avg_inference_time': round(self.avg_inference_time * 1000, 2) if self.avg_inference_time > 0 else 0,
            'total_detections': self.detection_count
        }
    
    def update_performance_stats(self, start_time: float, detection_count: int = 1):
        """更新性能统计信息
        
        Args:
            start_time: 检测开始时间
            detection_count: 本次检测的目标数量
        """
        elapsed = time.time() - start_time
        self.last_detection_time = time.time()
        self.detection_count += detection_count
        
        # 更新FPS（使用移动平均）
        current_fps = 1.0 / elapsed if elapsed > 0 else 0
        self.fps = self.fps * 0.9 + current_fps * 0.1
        
        # 更新平均推理时间
        self.total_inference_time += elapsed
        self.avg_inference_time = self.total_inference_time / self.detection_count
    
    def get_status(self) -> Dict[str, Any]:
        """获取检测引擎状态
        
        Returns:
            Dict: 引擎状态信息
        """
        return {
            'initialized': self.is_initialized,
            'device': self.device,
            'last_detection_time': self.last_detection_time,
            **self.get_performance_metrics()
        }
    
    def optimize(self):
        """优化检测引擎性能
        
        子类可以重写此方法以提供特定的优化
        """
        pass
    
    def release(self):
        """释放检测引擎资源"""
        pass